package com.siashan.toolkit.crypt.asymmetric;

import com.siashan.toolkit.crypt.CryptException;
import com.siashan.toolkit.crypt.KeyUtil;
import com.siashan.toolkit.crypt.SecureUtil;
import com.siashan.toolkit.crypt.binary.BCD;
import com.siashan.toolkit.crypt.util.StringUtils;

import javax.crypto.BadPaddingException;
import javax.crypto.Cipher;
import javax.crypto.IllegalBlockSizeException;
import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.nio.charset.Charset;
import java.nio.charset.StandardCharsets;
import java.security.*;
import java.security.spec.AlgorithmParameterSpec;

/**
 * 非对称加密算法
 *
 * <pre>
 * 1、签名：使用私钥加密，公钥解密。
 * 用于让所有公钥所有者验证私钥所有者的身份并且用来防止私钥所有者发布的内容被篡改，但是不用来保证内容不被他人获得。
 *
 * 2、加密：用公钥加密，私钥解密。
 * 用于向公钥所有者发布信息,这个信息可能被他人篡改,但是无法被他人获得。
 * </pre>
 *
 * @author sias
 * @since 1.0.7
 */
public class AsymmetricCrypto {

    private static  final Charset UTF8 = StandardCharsets.UTF_8;

    /**
     * Cipher负责完成加密或解密工作
     */
    protected Cipher cipher;

    protected String algorithm;

    /**
     * 加密的块大小
     */
    protected int encryptBlockSize = -1;
    /**
     * 解密的块大小
     */
    protected int decryptBlockSize = -1;

    /**
     * 算法参数
     */
    private AlgorithmParameterSpec algorithmParameterSpec;

    // ------------------------------------------------------------------ Constructor start

    /**
     * 构造，创建新的私钥公钥对
     *
     * @param algorithm {@link AsymmetricAlgorithm}
     */
    @SuppressWarnings("RedundantCast")
    public AsymmetricCrypto(AsymmetricAlgorithm algorithm) {
        this(algorithm.getValue());
    }

    /**
     * 构造，创建新的私钥公钥对
     *
     * @param algorithm 算法
     */
    @SuppressWarnings("RedundantCast")
    public AsymmetricCrypto(String algorithm) {
        this.algorithm = algorithm;
    }

    // ------------------------------------------------------------------ Constructor end

    /**
     * 获取加密块大小
     *
     * @return 加密块大小
     */
    public int getEncryptBlockSize() {
        return encryptBlockSize;
    }

    /**
     * 设置加密块大小
     *
     * @param encryptBlockSize 加密块大小
     */
    public void setEncryptBlockSize(int encryptBlockSize) {
        this.encryptBlockSize = encryptBlockSize;
    }

    /**
     * 获取解密块大小
     *
     * @return 解密块大小
     */
    public int getDecryptBlockSize() {
        return decryptBlockSize;
    }

    /**
     * 设置解密块大小
     *
     * @param decryptBlockSize 解密块大小
     */
    public void setDecryptBlockSize(int decryptBlockSize) {
        this.decryptBlockSize = decryptBlockSize;
    }

    /**
     * 获取{@link AlgorithmParameterSpec}<br>
     * 在某些算法中，需要特别的参数，例如在ECIES中，此处为IESParameterSpec
     *
     * @return {@link AlgorithmParameterSpec}
     * @since 5.4.3
     */
    public AlgorithmParameterSpec getAlgorithmParameterSpec() {
        return algorithmParameterSpec;
    }

    /**
     * 设置{@link AlgorithmParameterSpec}<br>
     * 在某些算法中，需要特别的参数，例如在ECIES中，此处为IESParameterSpec
     *
     * @param algorithmParameterSpec {@link AlgorithmParameterSpec}
     * @since 5.4.3
     */
    public void setAlgorithmParameterSpec(AlgorithmParameterSpec algorithmParameterSpec) {
        this.algorithmParameterSpec = algorithmParameterSpec;
    }

    public AsymmetricCrypto init(String algorithm, PrivateKey privateKey, PublicKey publicKey) {
        this.algorithm = algorithm;
        initCipher();
        return this;
    }

    // --------------------------------------------------------------------------------- Encrypt


    // --------------------------------------------------------------------------------- Getters and Setters


    /**
     * 获得加密或解密器
     *
     * @return 加密或解密
     * @since 5.4.3
     */
    public Cipher getCipher() {
        return cipher;
    }

    /**
     * 初始化{@link Cipher}，默认尝试加载BC库
     *
     * @since 4.5.2
     */
    protected void initCipher() {
        this.cipher = SecureUtil.createCipher(algorithm);
    }

    /**
     * 加密或解密
     *
     * @param data         被加密或解密的内容数据
     * @param maxBlockSize 最大块（分段）大小
     * @return 加密或解密后的数据
     * @throws IllegalBlockSizeException 分段异常
     * @throws BadPaddingException       padding错误异常
     * @throws IOException               IO异常，不会被触发
     */
    private byte[] doFinal(byte[] data, int maxBlockSize) throws IllegalBlockSizeException, BadPaddingException, IOException {
        // 模长
        final int dataLength = data.length;

        // 不足分段
        if (dataLength <= maxBlockSize) {
            return this.cipher.doFinal(data, 0, dataLength);
        }

        // 分段解密
        return doFinalWithBlock(data, maxBlockSize);
    }

    /**
     * 分段加密或解密
     *
     * @param data         数据
     * @param maxBlockSize 最大分段的段大小，不能为小于1
     * @return 加密或解密后的数据
     * @throws IllegalBlockSizeException 分段异常
     * @throws BadPaddingException       padding错误异常
     * @throws IOException               IO异常，不会被触发
     */
    private byte[] doFinalWithBlock(byte[] data, int maxBlockSize) throws IllegalBlockSizeException, BadPaddingException, IOException {
        final int dataLength = data.length;
        @SuppressWarnings("resource") final ByteArrayOutputStream out = new ByteArrayOutputStream();

        int offSet = 0;
        // 剩余长度
        int remainLength = dataLength;
        int blockSize;
        // 对数据分段处理
        while (remainLength > 0) {
            blockSize = Math.min(remainLength, maxBlockSize);
            out.write(cipher.doFinal(data, offSet, blockSize));

            offSet += blockSize;
            remainLength = dataLength - offSet;
        }

        return out.toByteArray();
    }

    /**
     * 初始化{@link Cipher}
     *
     * @param mode 模式，可选{@link Cipher#ENCRYPT_MODE}或者{@link Cipher#DECRYPT_MODE}
     * @param key  密钥
     * @throws InvalidAlgorithmParameterException 异常算法错误
     * @throws InvalidKeyException                异常KEY错误
     */
    private void initCipher(int mode, Key key) throws InvalidAlgorithmParameterException, InvalidKeyException {
        if (null != this.algorithmParameterSpec) {
            cipher.init(mode, key, this.algorithmParameterSpec);
        } else {
            cipher.init(mode, key);
        }
    }


    // --------------------------------------------------------------------------------- Encrypt
    /**
     * 加密
     *
     * @param data    被加密的bytes
     * @param key      私钥或公钥 {@link Key}
     * @return 加密后的bytes
     */
    public byte[] encrypt(byte[] data, Key key) {
        try {
            initCipher(Cipher.ENCRYPT_MODE, key);
            if (this.encryptBlockSize < 0) {
                // 在引入BC库情况下，自动获取块大小
                final int blockSize = this.cipher.getBlockSize();
                if (blockSize > 0) {
                    this.encryptBlockSize = blockSize;
                }
            }
            return doFinal(data, this.encryptBlockSize < 0 ? data.length : this.encryptBlockSize);
        } catch (Exception e) {
            throw new CryptException(e);
        }
    }

    /**
     * 加密
     *
     * @param data    被加密的bytes
     * @param keyByte 私钥或公钥
     * @param keyType 密钥类型
     * @return 加密后的bytes
     */
    public byte[] encrypt(byte[] data, byte[] keyByte, KeyType keyType) {
        return encrypt(data, KeyUtil.generateKey(algorithm, keyByte, keyType));
    }

    /**
     * 加密
     *
     * @param data    被加密的bytes
     * @param key     私钥或公钥
     * @param keyType 秘钥类型
     * @return Hex字符串
     */
    public byte[] encrypt(byte[] data, String key,KeyType keyType) {
        return encrypt(data,key.getBytes(UTF8),keyType);
    }




    /**
     * 加密
     *
     * @param data    被加密的字符串
     * @param key     私钥或公钥
     * @param keyType 秘钥类型
     * @return 加密后的bytes
     */
    public byte[] encrypt(String data, String key, KeyType keyType) {
        return encrypt(data.getBytes(UTF8),key.getBytes(UTF8), keyType);
    }



    /**
     * 分组加密
     *
     * @param data    数据
     * @param keyType 密钥类型
     * @return 加密后的密文
     */
    public String encryptBcd(String data,String key, KeyType keyType) {
        return BCD.bcdToStr(encrypt(data, key,keyType));
    }

    // --------------------------------------------------------------------------------- Decrypt

    /**
     * 解密
     *
     * @param data    被解密的bytes
     * @param key 私钥或公钥 {@link Key}
     * @return 解密后的bytes
     */
    public byte[] decrypt(byte[] data, Key key) {
        try {
            initCipher(Cipher.DECRYPT_MODE, key);
            if (this.decryptBlockSize < 0) {
                // 在引入BC库情况下，自动获取块大小
                final int blockSize = this.cipher.getBlockSize();
                if (blockSize > 0) {
                    this.decryptBlockSize = blockSize;
                }
            }

            return doFinal(data, this.decryptBlockSize < 0 ? data.length : this.decryptBlockSize);
        } catch (Exception e) {
            throw new CryptException(e);
        }
    }

    /**
     * 加密
     *
     * @param data    被加密的bytes
     * @param keyByte 私钥或公钥
     * @param keyType 密钥类型
     * @return 加密后的bytes
     */
    public byte[] decrypt(byte[] data, byte[] keyByte, KeyType keyType) {
        return decrypt(data, KeyUtil.generateKey(algorithm, keyByte, keyType));
    }

    /**
     * 从Hex或Base64字符串解密，编码为UTF-8格式
     *
     * @param data    Hex（16进制）或Base64字符串
     * @param key     秘钥
     * @param keyType 私钥或公钥 {@link KeyType}
     * @return 解密后的bytes
     */
    public byte[] decrypt(String data, String key,KeyType keyType) {
        return decrypt(data.getBytes(UTF8),key.getBytes(UTF8), keyType);
    }

    /**
     * 解密为字符串，密文需为Hex（16进制）或Base64字符串
     *
     * @param data    数据，Hex（16进制）或Base64字符串
     * @param keyType 密钥类型
     * @param charset 加密前编码
     * @return 解密后的密文
     * @since 4.5.2
     */
    public String decryptStr(String data,String key, KeyType keyType, Charset charset) {
        return StringUtils.newString(decrypt(data, key,keyType), charset);
    }

    /**
     * 解密为字符串，密文需为Hex（16进制）或Base64字符串
     *
     * @param data    数据，Hex（16进制）或Base64字符串
     * @param keyType 密钥类型
     * @return 解密后的密文
     * @since 4.5.2
     */
    public String decryptStr(String data,String key, KeyType keyType) {
        return decryptStr(data, key,keyType, UTF8);
    }

    /**
     * 解密BCD
     *
     * @param data    数据
     * @param keyType 密钥类型
     * @return 解密后的密文
     * @since 4.1.0
     */
    public byte[] decryptFromBcd(String data, String key,KeyType keyType) {
        return decryptFromBcd(data,key, keyType, UTF8);
    }

    /**
     * 分组解密
     *
     * @param data    数据
     * @param keyType 密钥类型
     * @param charset 加密前编码
     * @return 解密后的密文
     * @since 4.1.0
     */
    public byte[] decryptFromBcd(String data,String key, KeyType keyType, Charset charset) {
        final byte[] dataBytes = BCD.ascToBcd(data.getBytes(charset));
        return decrypt(dataBytes, key.getBytes(UTF8),keyType);
    }

    /**
     * 解密为字符串，密文需为BCD格式
     *
     * @param data    数据，BCD格式
     * @param keyType 密钥类型
     * @param charset 加密前编码
     * @return 解密后的密文
     * @since 4.5.2
     */
    public String decryptStrFromBcd(String data,String key, KeyType keyType, Charset charset) {
        return StringUtils.newString(decryptFromBcd(data,key, keyType, charset), charset);
    }

    /**
     * 解密为字符串，密文需为BCD格式，编码为UTF-8格式
     *
     * @param data    数据，BCD格式
     * @param keyType 密钥类型
     * @return 解密后的密文
     * @since 4.5.2
     */
    public String decryptStrFromBcd(String data,String key, KeyType keyType) {
        return decryptStrFromBcd(data, key,keyType, UTF8);
    }
}
